from collections import OrderedDict
from typing import Sequence

from paddle.fluid.io import DataLoader

from helm.static.engine.callback import Callback, CallbackList, CallOn


class State:

    def __init__(self, **kwargs):
        self.iteration = 0
        self.epoch = 0
        self.epoch_length = None
        self.max_epochs = None

        self.output = None
        self.batch = None
        self.metrics = {}
        self.dataloader = None
        self.seed = None

        for k, v in kwargs.items():
            setattr(self, k, v)


class Engine:
    _state_dict_all_req_keys = ("max_epochs", "epoch")

    def __init__(self, executor, step_fn, callbacks: Sequence[Callback], **kwargs):
        self.executor = executor
        self.step_fn = step_fn
        self.callbacks = CallbackList(callbacks)
        self.state = State()

        for k, v in kwargs.items():
            setattr(self, k, v)

    def call(self, callback: Callback):
        self.callbacks.append(callback)

    def call_on(self, event, f, freq=1):
        self.callbacks.append(CallOn(event, f, freq))

    def run(self, dataloader: DataLoader, num_epochs: int):
        state = self.state
        cbks = self.callbacks

        state.dataloader = dataloader
        state.epoch_length = len(dataloader)
        state.max_epochs = num_epochs

        cbks.on_begin(self)
        start_epoch = state.epoch + 1 if state.epoch != 0 else 0
        for epoch in range(start_epoch, num_epochs):
            state.epoch = epoch
            self._run_epoch()
        cbks.on_end(self)

    def _run_epoch(self):
        state = self.state
        cbks = self.callbacks

        cbks.on_epoch_begin(self)
        for iteration, batch in enumerate(state.dataloader):
            state.iteration = iteration
            state.batch = batch

            cbks.on_batch_begin(self)
            output = self.step_fn(self.executor, state.batch)
            state.output = output
            cbks.on_batch_end(self)
        cbks.on_epoch_end(self)

    def set_dict(self, state_dict):
        self.state = State(max_epochs=state_dict["max_epochs"], metrics={})

        self.state.epoch = state_dict["epoch"]

    def state_dict(self) -> OrderedDict:
        state = self.state
        return OrderedDict({
            "max_epochs": state.max_epochs,
            "epoch": state.epoch,
        })
